package com.feixiang.feixiangagent.chatmemory;
import com.esotericsoftware.kryo.Kryo;
import com.esotericsoftware.kryo.io.Input;
import com.esotericsoftware.kryo.io.Output;
import org.objenesis.strategy.StdInstantiatorStrategy;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.Message;
import java.io.File;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
/**
 * 基于文件持久化的对话记忆
 */
    public class FileBasedChatMemory implements ChatMemory {

        private final String BASE_DIR;
        private static final Kryo kryo = new Kryo();

        static {
            kryo.setRegistrationRequired(false);
            // 设置实例化策略
            kryo.setInstantiatorStrategy(new StdInstantiatorStrategy());
        }
        // 构造对象时，指定文件保存目录
        public FileBasedChatMemory(String dir) {
            this.BASE_DIR = dir;
            File baseDir = new File(dir);
            if (!baseDir.exists()) {
                baseDir.mkdirs();
            }
        }
        @Override
        public void add(String conversationId, List<Message> messages) {
            List<Message> conversationMessages = getOrCreateConversation(conversationId);
            conversationMessages.addAll(messages);
            saveConversation(conversationId, conversationMessages);
        }
        @Override
        public List<Message> get(String conversationId, int lastN) {
            List<Message> allMessages = getOrCreateConversation(conversationId);
            return allMessages.stream()
                    .skip(Math.max(0, allMessages.size() - lastN))
                    .toList();
        }
        @Override
        public void clear(String conversationId) {
            File file = getConversationFile(conversationId);
            if (file.exists()) {
                file.delete();
            }
        }
        private List<Message> getOrCreateConversation(String conversationId) {
            File file = getConversationFile(conversationId);
            List<Message> messages = new ArrayList<>();
            if (file.exists()) {
                try (Input input = new Input(new FileInputStream(file))) {
                    messages = kryo.readObject(input, ArrayList.class);
                } catch (IOException e) {
                    e.printStackTrace();
                }
            }
            return messages;
        }
        private void saveConversation(String conversationId, List<Message> messages) {
            File file = getConversationFile(conversationId);
            try (Output output = new Output(new FileOutputStream(file))) {
                kryo.writeObject(output, messages);
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
        private File getConversationFile(String conversationId) {
            return new File(BASE_DIR, conversationId + ".kryo");
        }
}
